#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
中心复合设计(CCD)实验矩阵生成器
支持自定义因子、水平范围和输出文件的实验设计矩阵生成
"""

import numpy as np
import pandas as pd
import os
import argparse
import json
from itertools import product

def get_default_factor_config():
    """
    获取默认的因子配置
    
    Returns:
        dict: 默认因子配置字典
    """
    return {
        'Sequencing_Depth': {
            'name': 'Sequencing_Depth',
            'low': 10,    # 10M reads
            'high': 50,   # 50M reads
            'unit': 'M',
            'type': 'continuous'
        },
        'Read_Length': {
            'name': 'Read_Length', 
            'low': 100,   # 100bp
            'high': 300,  # 300bp
            'unit': 'bp',
            'type': 'continuous'
        },
        'Complexity': {
            'name': 'Complexity',
            'low': -1,    # low complexity
            'high': 1,    # high complexity
            'levels': ['low', 'medium', 'high'],
            'type': 'categorical'
        },
        'Homology': {
            'name': 'Homology',
            'low': 5,     # 5% homology
            'high': 25,   # 25% homology
            'unit': '%',
            'type': 'continuous'
        },
        'Host_Proportion': {
            'name': 'Host_Proportion',
            'low': 70,    # 70% host
            'high': 95,   # 95% host
            'unit': '%',
            'type': 'continuous'
        }
    }

def generate_ccd_matrix(factor_config=None, n_center=6, alpha=1.5):
    """
    生成中心复合设计(CCD)实验矩阵
    
    Args:
        factor_config (dict): 因子配置字典，包含因子名称和水平范围
        n_center (int): 中心点重复次数，默认为6
        alpha (float): 轴向点位置参数，默认为1.5
    
    Returns:
        pandas.DataFrame: 包含实验设计的数据框
    """
    
    # 如果没有提供因子配置，使用默认配置
    if factor_config is None:
        factor_config = get_default_factor_config()
    
    n_factors = len(factor_config)
    
    # 手动生成中心复合设计矩阵
    # CCD包含三部分：因子点、轴向点、中心点
    
    # 1. 因子点 (2^k factorial points)
    factorial_points = []
    for combo in product([-1, 1], repeat=n_factors):
        factorial_points.append(list(combo))
    factorial_points = np.array(factorial_points)
    
    # 2. 轴向点 (2k axial points)
    # 使用较小的alpha值以避免超出合理范围
    alpha = 1.5  # 调整alpha值
    axial_points = []
    for i in range(n_factors):
        # 正轴向点
        point_pos = [0] * n_factors
        point_pos[i] = alpha
        axial_points.append(point_pos)
        # 负轴向点
        point_neg = [0] * n_factors
        point_neg[i] = -alpha
        axial_points.append(point_neg)
    axial_points = np.array(axial_points)
    
    # 3. 中心点 (center points)
    center_points = np.zeros((n_center, n_factors))
    
    # 合并所有点
    design_matrix = np.vstack([factorial_points, axial_points, center_points])
    
    # 将编码矩阵转换为实际参数值
    factor_names = list(factor_config.keys())
    decoded_matrix = np.zeros_like(design_matrix)
    
    for i, factor_name in enumerate(factor_names):
        factor = factor_config[factor_name]
        
        if factor.get('type') == 'categorical':
            # 对于分类变量，使用特殊处理
            levels = factor.get('levels', ['low', 'medium', 'high'])
            decoded_values = []
            for val in design_matrix[:, i]:
                if val <= -0.5:
                    decoded_values.append(0)  # 第一个水平
                elif val >= 0.5:
                    decoded_values.append(len(levels) - 1)  # 最后一个水平
                else:
                    decoded_values.append(len(levels) // 2)  # 中间水平
            decoded_matrix[:, i] = decoded_values
        else:
            print(factor_name)
            # 对于连续变量，线性缩放
            low_val = factor['low']
            high_val = factor['high']
            # 将设计矩阵范围映射到[low_val, high_val]
            # 首先将极值限制在合理范围内
            design_values = design_matrix[:, i]
            # 限制设计值在[-2, 2]范围内以避免极端值
            design_values = np.clip(design_values, -2, 2)
            # 将[-2, 2]范围映射到[low_val, high_val]
            decoded_matrix[:, i] = (design_values + 2) / 4 * (high_val - low_val) + low_val
    
    # 创建DataFrame
    df = pd.DataFrame(decoded_matrix, columns=factor_names)
    
    # 格式化数据
    df['Run_ID'] = [f'run{i+1:03d}' for i in range(len(df))]
    
    # 重新排列列顺序，将Run_ID放在第一列
    cols = ['Run_ID'] + factor_names
    df = df[cols]
    
    # 格式化特定列
    for factor_name in factor_names:
        factor = factor_config[factor_name]
        
        if factor.get('type') == 'categorical':
            # 处理分类变量
            levels = factor.get('levels', ['low', 'medium', 'high'])
            level_map = {i: f"{levels[i]}_{factor_name.lower()}" for i in range(len(levels))}
            df[factor_name] = df[factor_name].round(0).astype(int).map(level_map)
        else:
            # 处理连续变量
            unit = factor.get('unit', '')
            if unit:
                df[factor_name] = df[factor_name].round(1).astype(str) + unit
            else:
                df[factor_name] = df[factor_name].round(1).astype(str)
    
    return df

def save_to_csv(df, filename='ccd_experiment_matrix.csv'):
    """
    将实验矩阵保存为CSV文件
    
    Args:
        df (pandas.DataFrame): 实验设计矩阵
        filename (str): 输出文件名
    """
    output_path = os.path.join(os.getcwd(), filename)
    df.to_csv(output_path, index=False, encoding='utf-8')
    print(f"实验矩阵已保存到: {output_path}")
    print(f"总共生成 {len(df)} 个实验条件")
    return output_path

def parse_arguments():
    """
    解析命令行参数
    
    Returns:
        argparse.Namespace: 解析后的参数
    """
    parser = argparse.ArgumentParser(
        description='中心复合设计(CCD)实验矩阵生成器',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例用法:
  # 使用默认配置
  python generate_ccd_design.py
  
  # 指定输出文件
  python generate_ccd_design.py -o my_experiment.csv
  
  # 从JSON文件加载因子配置
  python generate_ccd_design.py -f factors.json
  
  # 设置中心点数量和alpha值
  python generate_ccd_design.py --center 8 --alpha 2.0
        """
    )
    
    parser.add_argument(
        '-f', '--factors',
        type=str,
        help='因子配置JSON文件路径'
    )
    
    parser.add_argument(
        '-o', '--output',
        type=str,
        default='ccd_experiment_matrix.csv',
        help='输出CSV文件名 (默认: ccd_experiment_matrix.csv)'
    )
    
    parser.add_argument(
        '--center',
        type=int,
        default=6,
        help='中心点重复次数 (默认: 6)'
    )
    
    parser.add_argument(
        '--alpha',
        type=float,
        default=1.5,
        help='轴向点位置参数 (默认: 1.5)'
    )
    
    parser.add_argument(
        '--preview',
        type=int,
        default=10,
        help='预览显示的行数 (默认: 10)'
    )
    
    return parser.parse_args()

def load_factor_config(config_file):
    """
    从JSON文件加载因子配置
    
    Args:
        config_file (str): JSON配置文件路径
    
    Returns:
        dict: 因子配置字典
    """
    try:
        with open(config_file, 'r', encoding='utf-8') as f:
            config = json.load(f)
        return config
    except FileNotFoundError:
        print(f"错误: 找不到配置文件 {config_file}")
        return None
    except json.JSONDecodeError as e:
        print(f"错误: JSON文件格式错误 - {e}")
        return None

def main():
    """
    主函数
    """
    args = parse_arguments()
    
    print("正在生成中心复合设计(CCD)实验矩阵...")
    
    try:
        # 加载因子配置
        if args.factors:
            factor_config = load_factor_config(args.factors)
            if factor_config is None:
                return
        else:
            factor_config = get_default_factor_config()
        
        # 生成实验矩阵
        experiment_matrix = generate_ccd_matrix(
            factor_config=factor_config,
            n_center=args.center,
            alpha=args.alpha
        )
        
        # 显示前几行
        print(f"\n实验矩阵预览 (前{args.preview}行):")
        print(experiment_matrix.head(args.preview))
        print(f"\n矩阵形状: {experiment_matrix.shape}")
        
        # 保存到CSV文件
        output_file = save_to_csv(experiment_matrix, args.output)
        
        # 显示统计信息
        print("\n实验设计统计信息:")
        print(f"- 因子数量: {len(factor_config)}")
        print(f"- 实验运行次数: {len(experiment_matrix)}")
        print(f"- 中心点重复次数: {args.center}")
        print(f"- 轴向点参数alpha: {args.alpha}")
        print(f"- 输出文件: {output_file}")
        
        print("\n各因子配置:")
        for factor_name, factor_info in factor_config.items():
            if factor_info.get('type') == 'categorical':
                levels = factor_info.get('levels', [])
                print(f"- {factor_name}: {'/'.join(levels)} (分类变量)")
            else:
                low = factor_info.get('low', 0)
                high = factor_info.get('high', 1)
                unit = factor_info.get('unit', '')
                print(f"- {factor_name}: {low}{unit} - {high}{unit} (连续变量)")
        
    except Exception as e:
        print(f"生成实验矩阵时发生错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()